#pragma once

#include "EquationSolver.h"

/**
 * @file RungeKutta.h
 * @author xingyifan
 * @date 2022-07-26 18:27
 *
 * @description: 定义时间网格和网格迭代器
 */

using TimeGrid = std::vector<Vector>;
using Gtimer = TimeGrid::iterator;

/**
 * @file RungeKutta.h
 * @author xingyifan
 * @date 2022-08-28 01:12
 * 
 * @description: 基类 RungeKutta 方法
 */

class RungeKutta
{
public:
    virtual Gtimer shift(FuncXt, Gtimer) = 0;
};


/**
 * @file RungeKutta.h
 * @author xingyifan
 * @date 2022-07-27 16:13
 *
 * @description: 经典 RungeKutta 方法
 */

class ClassicalRK : public RungeKutta
{
public:
    ClassicalRK(Real k) : t_(0), k_(k), A_(4, 0), b_(4, 0), c_(4, 0)
    {
        // RK 节点和权
        c_ = {0, 0.5, 0.5, 1};
        b_ = {1.0 / 6, 1.0 / 3, 1.0 / 3, 1.0 / 6};

        A_[1 * 4 + 0] = 0.5;
        A_[2 * 4 + 1] = 0.5;
        A_[3 * 4 + 2] = 1;
    }

    /**
     * @file RungeKutta.h
     * @author xingyifan
     * @date 2022-07-27 22:22
     *
     * @description: 显式方法，直接求解
     */

    Gtimer shift(FuncXt f, Gtimer U)
    {
        // 计算每个 y
        Vector y1 = f(*U, t_);
        Vector y2 = f(*U + y1 * A_[1 * 4 + 0] * k_, t_ + c_[1] * k_);
        Vector y3 = f(*U + y2 * A_[2 * 4 + 1] * k_, t_ + c_[2] * k_);
        Vector y4 = f(*U + y3 * A_[3 * 4 + 2] * k_, t_ + c_[3] * k_);

        // 加权求和
        *(U + 1) = *U + (y1 * b_[0] + y2 * b_[1] + y3 * b_[2] + y4 * b_[3]) * k_;

        // 时间推进
        t_ += k_;

        return U + 1;
    }

public:
    Matrix A_; // 系数矩阵
    Vector b_; // 加权
    Vector c_; // 节点

    Real t_; // 当前时间
    Real k_; // 时间步长
};

/**
 * @file RungeKutta.h
 * @author xingyifan
 * @date 2022-07-27 16:13
 *
 * @description: ESDI RK 方法
 */

class ESDIRK : public RungeKutta
{
public:
    ESDIRK(Real k) : s_(6), t_(0), k_(k), A_(6, 0), b_(6, 0), c_(6, 0)
    {
        // RK 节点和权
        c_ = {0, 0.5, 83.0 / 250, 31.0 / 50, 17.0 / 20, 1};
        b_ = {82889.0 / 524892, 0, 15625.0 / 83664, 69875.0 / 102672, -2260.0 / 8211, 0.25};

        A_[1 * 6 + 0] = 0.25;
        A_[1 * 6 + 1] = 0.25;

        A_[2 * 6 + 0] = 8611.0 / 62500;
        A_[2 * 6 + 1] = -1743.0 / 31250;
        A_[2 * 6 + 2] = 0.25;

        A_[3 * 6 + 0] = 5012029.0 / 34652500;
        A_[3 * 6 + 1] = -654441.0 / 2922500;
        A_[3 * 6 + 2] = 174375.0 / 388108;
        A_[3 * 6 + 3] = 0.25;

        A_[4 * 6 + 0] = 15267082809.0 / 155376265600;
        A_[4 * 6 + 1] = -71443401.0 / 120774400;
        A_[4 * 6 + 2] = 730878875.0 / 902184768;
        A_[4 * 6 + 3] = 2285395.0 / 8070912;
        A_[4 * 6 + 4] = 0.25;

        for (int i = 0; i < 6; i++)
        {
            A_[5 * 6 + i] = b_[i];
        }
    }

    /**
     * @file RungeKutta.h
     * @author xingyifan
     * @date 2022-07-27 22:41
     *
     * @description: 对角线隐式方法，对每个方程单独使用牛顿法求根
     */

    Gtimer shift(FuncXt f, Gtimer U)
    {
        // 存放 y 值，每个都初始化为 0
        std::vector<Vector> Y(s_, f(*U, t_));

        // 牛顿迭代器
        FSolve fS;

        // 对角线循环
        for (int i = 0; i < s_; i++)
        {
            // 构造隐式方程
            FuncX F = [=](Vector X) -> Vector
            {
                Vector ty = *U;
                for (int j = 0; j <= i; j++)
                {
                    ty = ty + Y[j] * A_[i * s_ + j] * k_;
                }
                return X - f(ty, t_ + c_[i] * k_);
            };

            // 迭代求根
            fS.solve(F, Y[i], true);
        }

        // 加权求和
        *(U + 1) = *U;
        for (int i = 0; i < s_; i++)
        {
            *(U + 1) = *(U + 1) + Y[i] * b_[i] * k_;
        }

        // 时间推进
        t_ += k_;

        return U + 1;
    }

    /**
     * @file RungeKutta.h
     * @author xingyifan
     * @date 2022-07-28 22:16
     *
     * @description: 对角隐式方程，当函数为线性形式：f(X,t) = AX + g(t) 时，可以通过求解线性系统推进
     */

    Gtimer shift(Matrix &A, Func g, Gtimer U)
    {
        // 矩阵规模
        int N = A.width();

        // 存放 phi 值，初始化为 *U
        std::vector<Vector> Phi(s_, *U);

        // 系数矩阵 B = (I - k a_ii A) 由于对角元相同，因此不重复计算
        Matrix B = A * 0.25 * k_;
        for (int i = 0; i < N; i++)
        {
            B[i * N + i] = 1 - B[i * N + i];
        }

        // 线性求解器
        LinearSolver lS;

        // 对角线循环
        for (int i = 1; i < s_; i++)
        {
            // b = U^n + k sum_j^(i-1) a_ij [A Phi[j] + g(t + cj * k)]
            for (int j = 0; j < i; j++)
            {
                Phi[i] = Phi[i] + (A * Phi[j] + g(t_ + c_[j] * k_)) * A_[i * s_ + j] * k_;
            }
            // 使用线性求解
            lS.solve(B, Phi[i], true);
        }

        // 加权求和
        *(U + 1) = *U;
        for (int i = 0; i < s_; i++)
        {
            *(U + 1) = *(U + 1) + (A * Phi[i] + g(t_ + c_[i] * k_)) * b_[i] * k_;
        }

        // 时间推进
        t_ += k_;

        return U + 1;
    }


public:
    Matrix A_; // 系数矩阵
    Vector b_; // 加权
    Vector c_; // 节点

    int s_;  // 级数
    Real t_; // 当前时间
    Real k_; // 时间步长
};

/**
 * @file RungeKutta.h
 * @author xingyifan
 * @date 2022-07-27 16:14
 *
 * @description: Gauss-Legendre 方法
 */

class GaussLegendre : public RungeKutta
{
public:
    GaussLegendre(int s, Real k) : s_(s), t_(0), k_(k), A_(s, 0), b_(s, 0), c_(s, 0)
    {
        // 获得 Gauss-Legendre 多项式
        Polynomial GL = Polynomial::GaussLegendre(s);

        // 使用牛顿法迭代求根
        FSolve fS;

        // 依次计算 s 个根，从 0 开始迭代，通常会得到从小到大的根
        for (int i = 0; i < s; i++)
        {
            // 构造函数
            FuncX F = [=](Vector x) -> Vector
            {
                return GL(x[0]);
            };

            // 迭代求根
            Vector x(0);
            fS.solve(F, x, true);
            c_[i] = x[0];

            // 除以因式 x - c[i] 进行降阶
            GL = GL / Polynomial({-c_[i], 1});
        }

        // 保险起见按照升序排列
        TensorOp::bubbleSort(c_);

        // 转换为向量
        std::vector<Real> c = c_;

        // 初始化权
        for (int i = 0; i < s; i++)
        {
            b_[i] = Polynomial::Lagrange(c, i).intergral(0, 1);
        }

        // 初始化系数矩阵
        for (int i = 0; i < s; i++)
        {
            for (int j = 0; j < s; j++)
            {
                A_[i * s + j] = Polynomial::Lagrange(c, j).intergral(0, c_[i]);
            }
        }
    }

    /**
     * @file RungeKutta.h
     * @author xingyifan
     * @date 2022-07-27 23:01
     *
     * @description: 纯隐式方法，使用向量牛顿迭代法求解
     */

    Gtimer shift(FuncXt f, Gtimer U)
    {
        int N = U->size();

        // 变量规模：共 s_ 个方程，每个方程是一个 N 元向量
        // Y = (y1 y2 ... ys)
        FuncX F = [=](Vector Y) -> Vector
        {
            // 结果容器
            Vector res(Y.size(), 0);

            // i ：第 i 个方程
            for (int i = 0; i < s_; i++)
            {
                // 获得初始向量
                Vector yi = *U;

                for (int j = 0; j < s_; j++)
                {
                    for (int n = 0; n < N; n++)
                    {
                        yi[n] += k_ * A_[i * s_ + j] * Y[j * N + n];
                    }
                }

                // 计算 yi
                yi = f(yi, t_ + k_ * c_[i]);

                // 得到第 i 方程
                for (int n = 0; n < N; n++)
                {
                    res[i * N + n] = Y[i * N + n] - yi[n];
                }
            }
            return res;
        };

        // 使用 f(U^n,t) 作为初值
        Vector Y0 = f(*U, t_);
        Vector Y(s_ * N, 0);

        for (int i = 0; i < s_; i++)
        {
            for (int n = 0; n < N; n++)
            {
                Y[i * N + n] = Y0[n];
            }
        }

        // 牛顿法迭代求根
        FSolve fS;
        fS.solve(F, Y, true);

        // 加权推进
        *(U + 1) = *U;
        for (int i = 0; i < s_; i++)
        {
            for (int n = 0; n < N; n++)
            {
                (*(U + 1))[n] += k_ * b_[i] * Y[i * N + n];
            }
        }

        // 时间推进
        t_ += k_;

        return U + 1;
    }

public:
    Matrix A_; // 系数矩阵
    Vector b_; // 加权
    Vector c_; // 节点

    int s_;  // 级数
    Real t_; // 当前时间
    Real k_; // 时间步长
};